# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import numpy as np

import paddle
from paddle.nn import ReflectionPad1d, ReflectionPad2d, ReflectionPad3d


class TestReflectionPad1d(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_ncl_format(self):
        # NCL: (1, 2, 3)
        data_np = np.array([[[1, 2, 3], [4, 5, 6]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = [1, 2]  # pad_left=1, pad_right=2
        my_pad = ReflectionPad1d(padding=pad, data_format="NCL")
        assert (
            my_pad.__repr__()
            == "ReflectionPad1D(padding=[1, 2], mode=reflect, value=0.0, data_format=NCL)"
        )
        result = my_pad(data)

        expected_np = np.array(
            [[[2.0, 1.0, 2.0, 3.0, 2.0, 1.0], [5.0, 4.0, 5.0, 6.0, 5.0, 4.0]]],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 2, 6])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_nlc_format(self):
        # NLC: (1, 3, 2)
        data_np = np.array([[[1, 4], [2, 5], [3, 6]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = [1, 2]  # pad_left=1, pad_right=2 (applies to L dim, size 3)
        my_pad = ReflectionPad1d(padding=pad, data_format="NLC")
        result = my_pad(data)

        expected_np = np.array(
            [
                [
                    [2.0, 5.0],
                    [1.0, 4.0],
                    [2.0, 5.0],
                    [3.0, 6.0],
                    [2.0, 5.0],
                    [1.0, 4.0],
                ]
            ],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 6, 2])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_int_padding(self):
        # NCL: (1, 2, 3)
        data_np = np.array([[[1, 2, 3], [4, 5, 6]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = 2  # (2, 2)
        my_pad = ReflectionPad1d(padding=pad, data_format="NCL")
        result = my_pad(data)

        expected_np = np.array(
            [
                [
                    [3.0, 2.0, 1.0, 2.0, 3.0, 2.0, 1.0],
                    [6.0, 5.0, 4.0, 5.0, 6.0, 5.0, 4.0],
                ]
            ],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 2, 7])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_padding_constraint(self):
        data = paddle.randn((1, 2, 3))
        with self.assertRaises(ValueError):
            my_pad = ReflectionPad1d(padding=3, data_format="NCL")
            my_pad(data)
        with self.assertRaises(ValueError):
            my_pad = ReflectionPad1d(padding=[1, 3], data_format="NCL")
            my_pad(data)


class TestReflectionPad2d(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_nchw_format(self):
        data_np = np.array([[[[1, 2, 3], [4, 5, 6]]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = [1, 0, 1, 1]
        my_pad = ReflectionPad2d(padding=pad, data_format="NCHW")
        assert (
            my_pad.__repr__()
            == "ReflectionPad2D(padding=[1, 0, 1, 1], mode=reflect, value=0.0, data_format=NCHW)"
        )
        result = my_pad(data)

        expected_np = np.array(
            [
                [
                    [
                        [5.0, 4.0, 5.0, 6.0],  # T=1 (R2_p)
                        [2.0, 1.0, 2.0, 3.0],  # R1_p
                        [5.0, 4.0, 5.0, 6.0],  # R2_p
                        [2.0, 1.0, 2.0, 3.0],  # B=1 (R1_p)
                    ]
                ]
            ],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 1, 4, 4])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_nhwc_format(self):
        data_np = np.array(
            [[[[1], [2], [3]], [[4], [5], [6]]]], dtype="float32"
        )
        data = paddle.to_tensor(data_np)
        pad = [1, 0, 1, 1]
        my_pad = ReflectionPad2d(padding=pad, data_format="NHWC")
        result = my_pad(data)

        expected_np = np.array(
            [
                [
                    [[5.0], [4.0], [5.0], [6.0]],  # T=1 (R2_p)
                    [[2.0], [1.0], [2.0], [3.0]],  # R1_p
                    [[5.0], [4.0], [5.0], [6.0]],  # R2_p
                    [[2.0], [1.0], [2.0], [3.0]],  # B=1 (R1_p)
                ]
            ],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 4, 4, 1])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_int_padding(self):
        data_np = np.array([[[[1, 2], [3, 4]]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = 1  # (1, 1, 1, 1)
        my_pad = ReflectionPad2d(padding=pad, data_format="NCHW")
        result = my_pad(data)

        paddings_np = ((0, 0), (0, 0), (1, 1), (1, 1))
        expected_np = np.pad(data_np, paddings_np, mode='reflect')

        self.assertEqual(result.shape, [1, 1, 4, 4])
        np.testing.assert_allclose(result.numpy(), expected_np)


class TestReflectionPad3d(unittest.TestCase):
    def setUp(self):
        paddle.disable_static()

    def test_ncdhw_format(self):
        data_np = np.array([[[[[1, 2, 3], [4, 5, 6]]]]], dtype="float32")
        data = paddle.to_tensor(data_np)

        pad = [1, 0, 1, 1, 0, 0]
        my_pad = ReflectionPad3d(padding=pad, data_format="NCDHW")
        assert (
            my_pad.__repr__()
            == "ReflectionPad3D(padding=[1, 0, 1, 1, 0, 0], mode=reflect, value=0.0, data_format=NCDHW)"
        )
        result = my_pad(data)

        expected_np = np.array(
            [
                [
                    [
                        [  # D=0
                            [5.0, 4.0, 5.0, 6.0],  # T=1
                            [2.0, 1.0, 2.0, 3.0],  # R1
                            [5.0, 4.0, 5.0, 6.0],  # R2
                            [2.0, 1.0, 2.0, 3.0],  # B=1
                        ]
                    ]
                ]
            ],
            dtype="float32",
        )

        self.assertEqual(result.shape, [1, 1, 1, 4, 4])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_int_padding_and_constraint(self):
        data_np = np.array(
            [[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]]], dtype="float32"
        )
        data = paddle.to_tensor(data_np)

        pad = 1  # (1, 1, 1, 1, 1, 1)
        my_pad = ReflectionPad3d(padding=pad, data_format="NCDHW")
        result = my_pad(data)

        paddings_np = ((0, 0), (0, 0), (1, 1), (1, 1), (1, 1))
        expected_np = np.pad(data_np, paddings_np, mode='reflect')

        self.assertEqual(result.shape, [1, 1, 4, 4, 4])
        np.testing.assert_allclose(result.numpy(), expected_np)

    def test_ndhwc_format(self):
        data_np = np.array(
            [[[[[1], [2]], [[3], [4]]], [[[5], [6]], [[7], [8]]]]],
            dtype="float32",
        )
        data = paddle.to_tensor(data_np)

        pad = 1  # (1, 1, 1, 1, 1, 1)
        my_pad = ReflectionPad3d(padding=pad, data_format="NDHWC")
        result = my_pad(data)

        paddings_np = ((0, 0), (1, 1), (1, 1), (1, 1), (0, 0))
        expected_np = np.pad(data_np, paddings_np, mode='reflect')

        self.assertEqual(result.shape, [1, 4, 4, 4, 1])
        np.testing.assert_allclose(result.numpy(), expected_np)


if __name__ == "__main__":
    unittest.main()
